from flask import Flask, request, jsonify
from ultralytics import YOLO

app = Flask(__name__)

# 加载模型
model = YOLO("D:/project/GitPro/model/major-events/backend/springboot-demo/src/main/java/com/qiu/demo/model/best.pt")

@app.route('/predict', methods=['POST'])
def predict():
    # 获取请求中的图片路径
    data = request.json
    image_path = data.get('image_path')

    if not image_path:
        return jsonify({'error': 'No image path provided'}), 400

    # 进行预测
    results = model([image_path], stream=True)

    res = []
    for result in results:
        if len(result.boxes.cls) > 0:
            for i in range(len(result.boxes.cls)):
                lb_id = int(result.boxes.cls[i].item())
                lb = result.names[lb_id]
                xsd = result.boxes.conf[i].item()
                zb = result.boxes.xyxy[i].tolist()
                res.append({'类别': lb, '相似度': xsd, '坐标': zb})

    return jsonify(res)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)